from rdkit import Chem
from rdkit.Chem import ChemicalFeatures
from rdkit import RDConfig
from rdkit.Chem.SaltRemover import SaltRemover
import numpy as np
from math import sqrt
import os

def remove_water(m):
  remover = SaltRemover(defnData="[O]")
  return remover.StripMol(m)

def get_feats(mol):
  fdefName = os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')
  factory = ChemicalFeatures.BuildFeatureFactory(fdefName)

  Donor_feat = factory.GetFeaturesForMol(mol, includeOnly="Donor")
  Acceptor_feat = factory.GetFeaturesForMol(mol, includeOnly="Acceptor")
  NegIonizable_feat = factory.GetFeaturesForMol(mol, includeOnly="NegIonizable")
  PosIonizable_feat = factory.GetFeaturesForMol(mol, includeOnly="PosIonizable")
  Aromatic_feat = factory.GetFeaturesForMol(mol, includeOnly="Aromatic")
  LumpedHydrophobe_feat = factory.GetFeaturesForMol(mol, includeOnly="LumpedHydrophobe")

  feats = Donor_feat + Acceptor_feat + NegIonizable_feat + PosIonizable_feat + Aromatic_feat + LumpedHydrophobe_feat
  return feats

def get_4_feats(mol):
    fdefName = os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')
    factory = ChemicalFeatures.BuildFeatureFactory(fdefName)

    Donor_feat = factory.GetFeaturesForMol(mol, includeOnly="Donor")
    Acceptor_feat = factory.GetFeaturesForMol(mol, includeOnly="Acceptor")
    Aromatic_feat = factory.GetFeaturesForMol(mol, includeOnly="Aromatic")
    # LumpedHydrophobe_feat = factory.GetFeaturesForMol(mol, includeOnly="LumpedHydrophobe")
    Hydrophobe_feat = factory.GetFeaturesForMol(mol, includeOnly="Hydrophobe")

    return Donor_feat + Acceptor_feat + Aromatic_feat + Hydrophobe_feat

def select_point(ligand_feats, protein_feats, distance):
    s_ligand_feats = []
    s_protein_feats = []

    c_l = np.zeros(len(ligand_feats))
    c_p = np.zeros(len(protein_feats))

    for i in range(len(ligand_feats)):
        l_pos = list(ligand_feats[i].GetPos())
        for j in range(len(protein_feats)):
            p_pos = list(protein_feats[j].GetPos())
            # print(euclidean_distance(l_pos, p_pos))
            if(sqrt(sum((a - b)**2 for a, b in zip(l_pos, p_pos))) <= distance):
                c_l[i] += 1
                c_p[j] += 1

        if c_l[i]:
            s_ligand_feats.append(ligand_feats[i])

    for i in range(len(protein_feats)):
        if c_p[i]:
            s_protein_feats.append(protein_feats[i])
            
    # print(len(protein_feats), c_l)
    return s_ligand_feats, s_protein_feats